import pymysql
import time
import pymysql.cursors
from contextlib import contextmanager
from pymongo import MongoClient
from .config import Config
from .models import db

# MongoDB客户端单例
_mongo_client = None

@contextmanager
def db_connection():
    """MySQL数据库连接上下文管理器"""
    # 配置校验
    required_fields = ['MYSQL_HOST', 'MYSQL_USER', 'MYSQL_PASSWORD', 'MYSQL_DB']
    for field in required_fields:
        if not hasattr(Config, field) or not getattr(Config, field):
            raise ValueError(f"Missing database configuration: {field}")

    conn = pymysql.connect(
        host=Config.MYSQL_HOST,
        port=Config.MYSQL_PORT,
        user=Config.MYSQL_USER,
        password=Config.MYSQL_PASSWORD,
        database=Config.MYSQL_DB,
        charset=Config.MYSQL_CHARSET,
        cursorclass=pymysql.cursors.DictCursor,
        connect_timeout=10,
        read_timeout=30,
        write_timeout=30
    )
    try:
        yield conn
    except pymysql.MySQLError as e:
        print(f"Database connection error: {str(e)}")
        raise
    finally:
        conn.close()

def get_mongo_client():
    """获取MongoDB客户端单例"""
    global _mongo_client
    if _mongo_client is None:
        _mongo_client = MongoClient(Config.MONGODB_URI)
    return _mongo_client

@contextmanager
def mongo_collection(collection_name):
    """MongoDB集合上下文管理器"""
    client = get_mongo_client()
    db = client[Config.MONGODB_DB]
    collection = db[collection_name]
    try:
        yield collection
    except Exception as e:
        print(f"MongoDB操作异常: {str(e)}")
        raise
    finally:
        # 这里不关闭连接，因为MongoClient会维护连接池
        pass

def initialize_mongodb():
    """初始化MongoDB数据库和集合"""
    try:
        client = get_mongo_client()
        mongo_db = client[Config.MONGODB_DB]
        
        print("MongoDB初始化完成")
    except Exception as e:
        print(f"MongoDB初始化失败: {str(e)}")
        raise

def init_app(app):
    """初始化SQLAlchemy"""
    db.init_app(app)
    with app.app_context():
        db.create_all()
    print("SQLAlchemy初始化完成")